Low-effort memory usage measurement in Rust
Me and my colleagues are doing Advent of Code 2023. One of them is doing it in C#, so he has access to a bunch of profiling tools and memory graphs and stuff. After day 5, he wanted to know what my memory usage was, because it was just that kind of problem.
I have a whole runner harness that runs my solutions (and provides them with the already-read input, because I don't wanna read the file by hand every time); it also measures runtime for each solution. I figured that'd be a good place to put the memory usage printout, too.
For actual measurement, after a brief search, I found peak_alloc. It's a crate that you can use as your allocator, and it keeps track of current and maximum memory allocation counts. It claims to have a "minimal overhead", so I could probably just use it, but I figured it's also time to get acquainted with cargo
s "features" feature, so I can toggle it on and off at will.
Using a custom allocator is easy
You literally just create it and slap an attribute on it:
#[global_allocator]
static PEAK_ALLOC: peak_alloc::PeakAlloc = peak_alloc::PeakAlloc;
Now every time memory is allocated, it uses code defined by the peak_alloc
crate to do so. Then, at the place I want to do the printout, I can literally just write:
println!("Peak memory usage: {:>3.3}mb", crate::PEAK_ALLOC.peak_usage_as_mb());
And that's pretty much it.
Turning it off and on again
Creating a feature flag for cargo
is actually trivial. If you mark a dependency as optional
, you're already done! That dependency is now also a feature flag:
[dependencies]
peak_alloc = { version = "0.2", optional = true }
With that modification, when I compile now, I get this error message:
error[E0433]: failed to resolve: use of undeclared crate or module peak_alloc
Which makes sense, it's optional now. If I do cargo run --features peak_alloc
instead, it works again. All that remains to do now, is slap some extra attributes on the sections that use peak_alloc
, so that it doesn't burn when I compile without the feature flag:
#[cfg(feature = "peak_alloc")]
#[global_allocator]
static PEAK_ALLOC: peak_alloc::PeakAlloc = peak_alloc::PeakAlloc;
#[cfg(feature = "peak_alloc")]
println!("Peak memory usage: {:>3.3}mb", crate::PEAK_ALLOC.peak_usage_as_mb());
And that's... pretty much it. Now it compiles with or without the feature flag, but prints peak memory usage with it. Nifty.
Giving the feature a better name
[features]
show_memory = ["peak_alloc"]
There, done. Now I can cargo run --features show_memory
instead. There's actually a small, subtle pitfall though: the dependency peak_alloc
and the feature flag peak_alloc
are actually distinct. If you read the relevant section in the Rust Book, it shows you that you can enable a dependency in a feature like so:
[feature]
show_memory = ["dep:peak_alloc"]
So now it would only turn on the dependency, but not the feature flag. Since my code is written with #[cfg(feature = "peak_alloc")]
, which checks against the feature flag, it would not show the memory usage. Yes, the dependency would be added, but the code that actually sets the allocator and prints the usage wouldn't actually be compiled in.